import os
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE' # 解决多个数据库冲突

from transformers import AutoTokenizer, AutoModel
import torch
import faiss
import numpy as np

def test_transformers():
    cache_folder = "D:/demo/gitee/python/models/text2vec-base-chinese"

    try:
        # 加载分词器
        tokenizer = AutoTokenizer.from_pretrained(
            "shibing624/text2vec-base-chinese",
            trust_remote_code=True,
            local_files_only=True,
            cache_dir=cache_folder
        )
        
        # 加载模型
        model = AutoModel.from_pretrained(
            "shibing624/text2vec-base-chinese",
            trust_remote_code=True,
            local_files_only=True,
            cache_dir=cache_folder
        )
        
        # 创建示例文本列表
        texts = ["你好", "今天天气真好", "我喜欢编程", "人工智能很有趣"]
        
        # 初始化向量数据库
        vector_dimension = 768  # 向量维度
        index = faiss.IndexFlatL2(vector_dimension)
        
        # 将多个文本转换为向量并添加到数据库
        all_embeddings = []
        for text in texts:
            # 文本向量化
            inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
            # 模型推理
            outputs = model(**inputs)
            # 获取文本向量
            embedding = outputs.last_hidden_state.mean(dim=1).detach().numpy()
            all_embeddings.append(embedding)
        
        # 将向量添加到FAISS索引
        all_embeddings = np.vstack(all_embeddings)
        index.add(all_embeddings)
        
        # 执行向量检索示例
        query_text = "编程很有意思"
        query_inputs = tokenizer(query_text, return_tensors="pt", padding=True, truncation=True)
        query_outputs = model(**query_inputs)
        query_embedding = query_outputs.last_hidden_state.mean(dim=1).detach().numpy()
        
        # 搜索最相似的向量
        D, I = index.search(query_embedding, k=2)  # 查找最相似的2个结果
        print(f"查询文本: {query_text}")
        print(f"最相似的文本索引: {I}")
        print(f"相似度距离: {D}")
        
    except Exception as e:
        print(f"加载模型出错：{e}")
    
    # from transformers import pipeline
    # classifier = pipeline("sentiment-analysis")
    # result = classifier(text)
    # print(result)


if __name__ == "__main__":
    test_transformers()